#!/usr/bin/env python3
import argparse
from typing import Dict, Optional, Tuple

import onnx
from onnx import helper, numpy_helper, TensorProto, shape_inference, ModelProto, NodeProto, GraphProto, ValueInfoProto
from rich import print
from rich.table import Table
import numpy as np


def _get_io_shapes_types(model: ModelProto) -> Dict[str, Tuple[Optional[str], Optional[Tuple[Optional[int], ...]]]]:
    """Return a dict mapping value name -> (dtype_str, shape_tuple) after shape inference."""
    inferred = shape_inference.infer_shapes(model)
    value_map: Dict[str, Tuple[Optional[str], Optional[Tuple[Optional[int], ...]]]] = {}

    def collect_vi(vi: ValueInfoProto):
        t = vi.type.tensor_type
        dtype = None
        if t.elem_type:
            dtype = TensorProto.DataType.Name(t.elem_type)
        shape = None
        if t.HasField("shape"):
            dims = []
            for d in t.shape.dim:
                if d.HasField("dim_value"):
                    dims.append(int(d.dim_value))
                elif d.HasField("dim_param"):
                    # symbolic
                    dims.append(None)
                else:
                    dims.append(None)
            shape = tuple(dims)
        value_map[vi.name] = (dtype, shape)

    for vi in list(inferred.graph.value_info) + list(inferred.graph.input) + list(inferred.graph.output):
        collect_vi(vi)
    return value_map


def print_model_io(model_path: str):
    model = onnx.load(model_path)
    io_map = _get_io_shapes_types(model)

    # Print inputs
    print("[bold cyan]Model Inputs[/bold cyan]")
    t = Table("name", "dtype", "shape")
    for inp in model.graph.input:
        dtype, shape = io_map.get(inp.name, (None, None))
        t.add_row(inp.name, str(dtype), str(shape))
    print(t)

    # Print outputs
    print("[bold cyan]Model Outputs[/bold cyan]")
    t = Table("name", "dtype", "shape")
    for out in model.graph.output:
        dtype, shape = io_map.get(out.name, (None, None))
        t.add_row(out.name, str(dtype), str(shape))
    print(t)

    # Print nodes
    print("[bold cyan]Nodes[/bold cyan]")
    nt = Table("idx", "op_type", "name", "inputs (shape)", "outputs (shape)")
    for idx, node in enumerate(model.graph.node):
        in_desc = []
        for n in node.input:
            dtype, shape = io_map.get(n, (None, None))
            in_desc.append(f"{n} ({dtype},{shape})")
        out_desc = []
        for n in node.output:
            dtype, shape = io_map.get(n, (None, None))
            out_desc.append(f"{n} ({dtype},{shape})")
        nt.add_row(str(idx), node.op_type, node.name or "-", "\n".join(in_desc), "\n".join(out_desc))
    print(nt)


def insert_input_transpose_nhwc(model: ModelProto) -> ModelProto:
    g = model.graph
    if len(g.input) != 1:
        raise ValueError(f"Expected a single model input, found {len(g.input)}")

    orig_input = g.input[0]
    orig_name = orig_input.name

    # Read original shape (expect NCHW)
    t = orig_input.type.tensor_type
    if not t.HasField("shape") or len(t.shape.dim) != 4:
        raise ValueError("Model input is not 4D")

    # Get dims if static, else fall back to defaults
    dims = []
    for d in t.shape.dim:
        if d.HasField("dim_value"):
            dims.append(int(d.dim_value))
        else:
            dims.append(None)

    # Default to YOLO common shape if unknown
    n, c, h, w = (dims + [None, None, None, None])[:4]
    if n is None:
        n = 1
    if c is None:
        c = 3
    if h is None:
        h = 640
    if w is None:
        w = 640

    # Create a new graph input with NHWC shape but SAME name as original
    # Then we will insert a Transpose right after input and route its output to original consumers.
    nhwc_dims = [n, h, w, c]

    # Update the existing input's shape to NHWC
    new_input = helper.make_tensor_value_info(orig_name, t.elem_type or TensorProto.FLOAT, nhwc_dims)

    # Build transpose node: NHWC -> NCHW (perm [0,3,1,2])
    transpose_out = orig_name + "__nchw"
    transpose_node = helper.make_node(
        "Transpose",
        inputs=[orig_name],
        outputs=[transpose_out],
        perm=[0, 3, 1, 2],
        name="Input_NHWC_to_NCHW",
    )

    # Replace all occurrences of original input name in node inputs with transpose_out
    new_nodes: list[NodeProto] = []
    for node in g.node:
        new_inputs = [transpose_out if i == orig_name else i for i in node.input]
        new_node = helper.make_node(
            node.op_type,
            inputs=new_inputs,
            outputs=list(node.output),
            name=node.name,
        )
        # Copy attributes
        for attr in node.attribute:
            new_node.attribute.extend([attr])
        new_nodes.append(new_node)

    # Rebuild graph
    new_graph = helper.make_graph(
        nodes=[transpose_node] + new_nodes,
        name=g.name or "graph",
        inputs=[new_input],
        outputs=list(g.output),
        initializer=list(g.initializer),
        value_info=list(g.value_info),
    )

    new_model = helper.make_model(new_graph, producer_name="yolo-transpose-tool")
    new_model.ir_version = model.ir_version
    # Properly copy opset imports (repeated field doesn't support slice assignment)
    del new_model.opset_import[:]
    for imp in model.opset_import:
        ops = onnx.OperatorSetIdProto()
        ops.domain = imp.domain
        ops.version = imp.version
        new_model.opset_import.extend([ops])
    onnx.checker.check_model(new_model)
    return new_model


def main():
    parser = argparse.ArgumentParser(description="Inspect ONNX model and insert an input Transpose to accept NHWC input.")
    parser.add_argument("model", type=str, help="Path to the original ONNX model (expects NCHW input)")
    parser.add_argument("--out", type=str, default=None, help="Path to save the transformed model (NHWC input)")
    parser.add_argument("--print", dest="do_print", action="store_true", help="Print node I/O information")
    args = parser.parse_args()

    print(f"[bold]Loading[/bold] {args.model}")
    model = onnx.load(args.model)

    if args.do_print:
        print_model_io(args.model)

    print("[bold]Transforming model to accept NHWC input by inserting Transpose at input...[/bold]")
    new_model = insert_input_transpose_nhwc(model)

    out_path = args.out or (args.model.rsplit(".", 1)[0] + "_nhwc.onnx")
    onnx.save(new_model, out_path)
    print(f"[green]Saved transformed model to[/green] {out_path}")

    if args.do_print:
        print("[bold]Transformed model I/O:[/bold]")
        print_model_io(out_path)


if __name__ == "__main__":
    main()
